package org.finesys.ai.core;

import static org.finesys.ai.core.HeaderNames.ORGANIZATION;
import static org.finesys.ai.core.HeaderNames.USER_AGENT;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.finesys.ai.core.embedding.EmbeddingRequest;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import okhttp3.Cache;
import okhttp3.OkHttpClient;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

/**
 * 内嵌客户端
 */

public class EmbeddingClient extends OpenAiClient {
    private static final Logger log = LoggerFactory.getLogger(EmbeddingClient.class);

    private final String baseUrl;

    private final String model;

    private final String apiVersion;

    private final OkHttpClient okHttpClient;

    private final OpenAiApi openAiApi;

    public EmbeddingClient(String apiKey) {
        this(new Builder().openAiApiKey(apiKey));
    }

    private EmbeddingClient(Builder builder) {
        this.baseUrl = builder.baseUrl;
        this.apiVersion = builder.apiVersion;
        this.model = builder.model;
        OkHttpClient.Builder okhttpClientBuilder = new OkHttpClient.Builder().callTimeout(builder.callTimeout)
                .connectTimeout(builder.connectTimeout).readTimeout(builder.readTimeout).writeTimeout(builder.writeTimeout);
        if (builder.dispatcher != null) {
            okhttpClientBuilder.dispatcher(builder.dispatcher);
        }
        if (StringUtils.hasText(builder.openAiApiKey)) {
            okhttpClientBuilder.addInterceptor(new AuthorizationHeaderInjector(builder.openAiApiKey));
        } else {
            okhttpClientBuilder.addInterceptor(new ApiKeyHeaderInjector(builder.azureApiKey));
        }
        Map<String, String> headers = new HashMap<>();
        if (StringUtils.hasText(builder.organizationId)) {
            headers.put(ORGANIZATION, builder.organizationId);
        }

        if (StringUtils.hasText(builder.userAgent)) {
            headers.put(USER_AGENT, builder.userAgent);
        }
        if (builder.customHeaders != null) {
            headers.putAll(builder.customHeaders);
        }
        if (!headers.isEmpty()) {
            okhttpClientBuilder.addInterceptor(new GenericHeaderInjector(headers));
        }
        if (builder.proxy != null) {
            okhttpClientBuilder.proxy(builder.proxy);
        }
        if (builder.logRequests) {
            okhttpClientBuilder.addInterceptor(new RequestLoggingInterceptor(builder.logLevel));
        }
        if (builder.logResponses) {
            okhttpClientBuilder.addInterceptor(new ResponseLoggingInterceptor(builder.logLevel));
        }
        this.okHttpClient = okhttpClientBuilder.build();
        Retrofit.Builder retrofitBuilder = new Retrofit.Builder().baseUrl(builder.baseUrl).client(okHttpClient);
        if (builder.persistTo != null) {
            retrofitBuilder.addConverterFactory(new PersistorConverterFactory(builder.persistTo));
        }
        retrofitBuilder.addConverterFactory(JacksonConverterFactory.create(Json.OBJECT_MAPPER));
        this.openAiApi = retrofitBuilder.build().create(OpenAiApi.class);
    }

    @Override
    public void shutdown() {
        okHttpClient.dispatcher().executorService().shutdown();
        okHttpClient.connectionPool().evictAll();
        Cache cache = okHttpClient.cache();
        if (cache != null) {
            try {
                cache.close();
            } catch (IOException e) {
                log.error("Failed to close cache {}", e.getLocalizedMessage());
            }
        }
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder extends OpenAiClient.Builder<EmbeddingClient, Builder> {
        /**
         * 构建并返回一个 EmbeddingClient 实例。
         *
         * @return 返回一个 EmbeddingClient 实例
         */
        @Override
        public EmbeddingClient build() {
            return new EmbeddingClient(this);
        }
    }


    /**
     * 根据给定的EmbeddingRequest生成嵌入表示
     *
     * @param request 请求对象，包含了生成嵌入表示所需的信息
     * @return 嵌入表示的结果，封装在EmbeddingResponse对象中
     */
    public EmbeddingResponse embed(EmbeddingRequest request) {
        return this.embedding(new OpenAiClientContext(), request).execute();
    }

    /**
     * 将输入文本嵌入到浮点数列表中。
     *
     * @param input 输入的文本
     * @return 包含嵌入结果的浮点数列表
     */
    public List<Double> embed(String input) {
        return this.embedding(new OpenAiClientContext(), input).execute();
    }

    @Override
    public SyncOrAsync<EmbeddingResponse> embedding(OpenAiClientContext clientContext, EmbeddingRequest request) {
        if (StringUtils.hasText(request.model())) {
            request.model();
        }
        return new RequestExecutor<>(openAiApi.embeddings(clientContext.headers(), request, apiVersion), r -> r);
    }

    @Override
    public SyncOrAsync<List<Double>> embedding(OpenAiClientContext clientContext, String input) {
        EmbeddingRequest.Builder builder = EmbeddingRequest.builder().input(input);
        if (StringUtils.hasText(this.model)) {
            builder.model(this.model);
        }

        return new RequestExecutor<>(openAiApi.embeddings(clientContext.headers(), builder.build(), apiVersion), EmbeddingResponse::embedding);
    }
}
